#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
KWDB SDK 功能测试和性能测试脚本
使用test-datasets中的NOAA数据进行测试
"""

import os
import sys
import time
import subprocess
import threading
from concurrent.futures import ThreadPoolExecutor
import json

class KWDBTester:
    def __init__(self):
        self.java_sdk_path = "/root/KWDB-sdk/java-sdk"
        self.cpp_sdk_path = "/root/KWDB-sdk/cpp-sdk"
        self.server_path = "/root/KWDB-sdk/server"
        self.test_data_path = "/root/KWDB-sdk/test-datasets/NOAA_data.txt"
        self.results = {}
        
    def check_dependencies(self):
        """检查依赖环境"""
        print("=== 检查依赖环境 ===")
        
        # 检查Java环境
        try:
            result = subprocess.run(['java', '-version'], capture_output=True, text=True)
            if result.returncode == 0:
                print("✓ Java环境正常")
            else:
                print("✗ Java环境异常")
                return False
        except FileNotFoundError:
            print("✗ 未找到Java环境")
            return False
            
        # 检查Maven
        try:
            result = subprocess.run(['mvn', '-version'], capture_output=True, text=True)
            if result.returncode == 0:
                print("✓ Maven环境正常")
            else:
                print("✗ Maven环境异常")
                return False
        except FileNotFoundError:
            print("✗ 未找到Maven环境")
            return False
            
        # 检查CMake
        try:
            result = subprocess.run(['cmake', '--version'], capture_output=True, text=True)
            if result.returncode == 0:
                print("✓ CMake环境正常")
            else:
                print("✗ CMake环境异常")
                return False
        except FileNotFoundError:
            print("✗ 未找到CMake环境")
            return False
            
        return True
    
    def build_java_sdk(self):
        """构建Java SDK"""
        print("\n=== 构建Java SDK ===")
        try:
            os.chdir(self.java_sdk_path)
            result = subprocess.run(['mvn', 'clean', 'compile', 'package'], 
                                 capture_output=True, text=True, timeout=300)
            if result.returncode == 0:
                print("✓ Java SDK构建成功")
                return True
            else:
                print(f"✗ Java SDK构建失败: {result.stderr}")
                return False
        except subprocess.TimeoutExpired:
            print("✗ Java SDK构建超时")
            return False
        except Exception as e:
            print(f"✗ Java SDK构建异常: {e}")
            return False
    
    def build_cpp_sdk(self):
        """构建C++ SDK"""
        print("\n=== 构建C++ SDK ===")
        try:
            os.chdir(self.cpp_sdk_path)
            
            # 创建构建目录
            if not os.path.exists('build'):
                os.makedirs('build')
            os.chdir('build')
            
            # 配置CMake
            result = subprocess.run(['cmake', '..'], capture_output=True, text=True, timeout=60)
            if result.returncode != 0:
                print(f"✗ CMake配置失败: {result.stderr}")
                return False
            
            # 编译
            result = subprocess.run(['make', '-j4'], capture_output=True, text=True, timeout=300)
            if result.returncode == 0:
                print("✓ C++ SDK构建成功")
                return True
            else:
                print(f"✗ C++ SDK构建失败: {result.stderr}")
                return False
        except subprocess.TimeoutExpired:
            print("✗ C++ SDK构建超时")
            return False
        except Exception as e:
            print(f"✗ C++ SDK构建异常: {e}")
            return False
    
    def build_server(self):
        """构建服务器"""
        print("\n=== 构建服务器 ===")
        try:
            os.chdir(self.server_path)
            result = subprocess.run(['mvn', 'clean', 'compile', 'package'], 
                                 capture_output=True, text=True, timeout=300)
            if result.returncode == 0:
                print("✓ 服务器构建成功")
                return True
            else:
                print(f"✗ 服务器构建失败: {result.stderr}")
                return False
        except subprocess.TimeoutExpired:
            print("✗ 服务器构建超时")
            return False
        except Exception as e:
            print(f"✗ 服务器构建异常: {e}")
            return False
    
    def start_server(self):
        """启动服务器"""
        print("\n=== 启动服务器 ===")
        try:
            os.chdir(self.server_path)
            # 启动服务器（后台运行）
            self.server_process = subprocess.Popen(
                ['java', '-jar', 'target/kwdb-server-1.0.0.jar', '--port', '9090'],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            
            # 等待服务器启动
            time.sleep(5)
            
            # 检查服务器是否启动成功
            if self.server_process.poll() is None:
                print("✓ 服务器启动成功")
                return True
            else:
                print("✗ 服务器启动失败")
                return False
        except Exception as e:
            print(f"✗ 服务器启动异常: {e}")
            return False
    
    def stop_server(self):
        """停止服务器"""
        if hasattr(self, 'server_process') and self.server_process:
            self.server_process.terminate()
            self.server_process.wait()
            print("✓ 服务器已停止")
    
    def parse_test_data(self):
        """解析测试数据"""
        print("\n=== 解析测试数据 ===")
        data_points = []
        
        try:
            with open(self.test_data_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            for line in lines:
                line = line.strip()
                if line and not line.startswith('#') and not line.startswith('CREATE'):
                    # 解析Line Protocol格式
                    parts = line.split(' ')
                    if len(parts) >= 3:
                        measurement_tags = parts[0]
                        fields = parts[1]
                        timestamp = parts[2]
                        
                        # 分离measurement和tags
                        if ',' in measurement_tags:
                            measurement, tags_part = measurement_tags.split(',', 1)
                        else:
                            measurement = measurement_tags
                            tags_part = ""
                        
                        data_points.append({
                            'measurement': measurement,
                            'tags': tags_part,
                            'fields': fields,
                            'timestamp': timestamp,
                            'line_protocol': line
                        })
            
            print(f"✓ 解析了 {len(data_points)} 条测试数据")
            return data_points
            
        except Exception as e:
            print(f"✗ 解析测试数据失败: {e}")
            return []
    
    def test_java_sdk(self, data_points):
        """测试Java SDK"""
        print("\n=== 测试Java SDK ===")
        
        # 创建测试程序
        test_code = '''
import com.kwdb.sdk.KWDBClient;
import com.kwdb.sdk.config.ClientConfig;
import com.kwdb.sdk.data.DataPoint;
import com.kwdb.sdk.exception.KWDBException;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

public class JavaSDKTest {
    public static void main(String[] args) {
        try {
            // 创建客户端配置
            ClientConfig config = new ClientConfig("localhost", 9090)
                    .setBatchSize(1000)
                    .setBatchTimeout(100)
                    .setEnableCompression(true);
            
            // 创建客户端
            KWDBClient client = new KWDBClient(config);
            client.connect();
            
            System.out.println("Java SDK连接成功");
            
            // 测试单条写入
            DataPoint dataPoint = new DataPoint("test_measurement")
                    .addTag("location", "test")
                    .addField("value", 1.0);
            
            long startTime = System.currentTimeMillis();
            client.write(dataPoint);
            long endTime = System.currentTimeMillis();
            
            System.out.println("单条写入耗时: " + (endTime - startTime) + "ms");
            
            // 测试批量写入
            List<DataPoint> batchData = new ArrayList<>();
            for (int i = 0; i < 100; i++) {
                DataPoint dp = new DataPoint("batch_test")
                        .addTag("id", String.valueOf(i))
                        .addField("value", i * 1.0);
                batchData.add(dp);
            }
            
            startTime = System.currentTimeMillis();
            client.writeBatch(batchData);
            endTime = System.currentTimeMillis();
            
            System.out.println("批量写入100条耗时: " + (endTime - startTime) + "ms");
            
            // 测试查询
            String sql = "SELECT * FROM test_measurement LIMIT 10";
            var result = client.query(sql);
            System.out.println("查询结果行数: " + result.getRows().size());
            
            client.disconnect();
            System.out.println("Java SDK测试完成");
            
        } catch (Exception e) {
            System.err.println("Java SDK测试失败: " + e.getMessage());
            e.printStackTrace();
        }
    }
}
'''
        
        # 写入测试文件
        test_file = os.path.join(self.java_sdk_path, "JavaSDKTest.java")
        with open(test_file, 'w', encoding='utf-8') as f:
            f.write(test_code)
        
        try:
            # 编译测试程序
            os.chdir(self.java_sdk_path)
            result = subprocess.run(['javac', '-cp', 'target/classes:target/dependency/*', 'JavaSDKTest.java'], 
                                 capture_output=True, text=True, timeout=60)
            
            if result.returncode != 0:
                print(f"✗ 编译测试程序失败: {result.stderr}")
                return False
            
            # 运行测试程序
            result = subprocess.run(['java', '-cp', '.:target/classes:target/dependency/*', 'JavaSDKTest'], 
                                 capture_output=True, text=True, timeout=30)
            
            if result.returncode == 0:
                print("✓ Java SDK测试成功")
                print(result.stdout)
                return True
            else:
                print(f"✗ Java SDK测试失败: {result.stderr}")
                return False
                
        except Exception as e:
            print(f"✗ Java SDK测试异常: {e}")
            return False
    
    def performance_test(self):
        """性能测试"""
        print("\n=== 性能测试 ===")
        
        # 测试写入性能
        print("测试写入性能...")
        
        # 创建性能测试程序
        perf_test_code = '''
import com.kwdb.sdk.KWDBClient;
import com.kwdb.sdk.config.ClientConfig;
import com.kwdb.sdk.data.DataPoint;
import com.kwdb.sdk.exception.KWDBException;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class PerformanceTest {
    public static void main(String[] args) {
        try {
            ClientConfig config = new ClientConfig("localhost", 9090)
                    .setBatchSize(1000)
                    .setBatchTimeout(100)
                    .setEnableCompression(true)
                    .setThreadPoolSize(20);
            
            KWDBClient client = new KWDBClient(config);
            client.connect();
            
            // 测试单线程写入性能
            int recordCount = 10000;
            List<DataPoint> dataPoints = new ArrayList<>();
            
            for (int i = 0; i < recordCount; i++) {
                DataPoint dp = new DataPoint("perf_test")
                        .addTag("thread", "single")
                        .addTag("id", String.valueOf(i))
                        .addField("value", Math.random() * 100)
                        .addField("timestamp", System.currentTimeMillis());
                dataPoints.add(dp);
            }
            
            long startTime = System.currentTimeMillis();
            client.writeBatch(dataPoints);
            long endTime = System.currentTimeMillis();
            
            long duration = endTime - startTime;
            double throughput = (double) recordCount / duration * 1000; // 记录/秒
            
            System.out.println("单线程写入性能:");
            System.out.println("记录数: " + recordCount);
            System.out.println("耗时: " + duration + "ms");
            System.out.println("吞吐量: " + String.format("%.2f", throughput) + " 记录/秒");
            
            // 测试多线程写入性能
            int threadCount = 10;
            int recordsPerThread = 1000;
            ExecutorService executor = Executors.newFixedThreadPool(threadCount);
            
            List<CompletableFuture<Void>> futures = new ArrayList<>();
            
            startTime = System.currentTimeMillis();
            
            for (int t = 0; t < threadCount; t++) {
                final int threadId = t;
                CompletableFuture<Void> future = CompletableFuture.runAsync(() -> {
                    try {
                        List<DataPoint> threadData = new ArrayList<>();
                        for (int i = 0; i < recordsPerThread; i++) {
                            DataPoint dp = new DataPoint("perf_test")
                                    .addTag("thread", "multi")
                                    .addTag("thread_id", String.valueOf(threadId))
                                    .addTag("id", String.valueOf(i))
                                    .addField("value", Math.random() * 100)
                                    .addField("timestamp", System.currentTimeMillis());
                            threadData.add(dp);
                        }
                        client.writeBatch(threadData);
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }, executor);
                futures.add(future);
            }
            
            // 等待所有线程完成
            CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
            
            endTime = System.currentTimeMillis();
            duration = endTime - startTime;
            int totalRecords = threadCount * recordsPerThread;
            throughput = (double) totalRecords / duration * 1000;
            
            System.out.println("\\n多线程写入性能:");
            System.out.println("线程数: " + threadCount);
            System.out.println("总记录数: " + totalRecords);
            System.out.println("耗时: " + duration + "ms");
            System.out.println("吞吐量: " + String.format("%.2f", throughput) + " 记录/秒");
            
            executor.shutdown();
            client.disconnect();
            
        } catch (Exception e) {
            System.err.println("性能测试失败: " + e.getMessage());
            e.printStackTrace();
        }
    }
}
'''
        
        # 写入性能测试文件
        perf_file = os.path.join(self.java_sdk_path, "PerformanceTest.java")
        with open(perf_file, 'w', encoding='utf-8') as f:
            f.write(perf_test_code)
        
        try:
            # 编译性能测试程序
            os.chdir(self.java_sdk_path)
            result = subprocess.run(['javac', '-cp', 'target/classes:target/dependency/*', 'PerformanceTest.java'], 
                                 capture_output=True, text=True, timeout=60)
            
            if result.returncode != 0:
                print(f"✗ 编译性能测试程序失败: {result.stderr}")
                return False
            
            # 运行性能测试
            result = subprocess.run(['java', '-cp', '.:target/classes:target/dependency/*', 'PerformanceTest'], 
                                 capture_output=True, text=True, timeout=120)
            
            if result.returncode == 0:
                print("✓ 性能测试完成")
                print(result.stdout)
                return True
            else:
                print(f"✗ 性能测试失败: {result.stderr}")
                return False
                
        except Exception as e:
            print(f"✗ 性能测试异常: {e}")
            return False
    
    def run_all_tests(self):
        """运行所有测试"""
        print("开始KWDB SDK功能测试和性能测试")
        print("=" * 50)
        
        # 检查依赖
        if not self.check_dependencies():
            print("依赖检查失败，退出测试")
            return False
        
        # 构建项目
        if not self.build_java_sdk():
            print("Java SDK构建失败，退出测试")
            return False
        
        if not self.build_server():
            print("服务器构建失败，退出测试")
            return False
        
        # 启动服务器
        if not self.start_server():
            print("服务器启动失败，退出测试")
            return False
        
        try:
            # 解析测试数据
            data_points = self.parse_test_data()
            
            # 测试Java SDK
            if not self.test_java_sdk(data_points):
                print("Java SDK测试失败")
                return False
            
            # 性能测试
            if not self.performance_test():
                print("性能测试失败")
                return False
            
            print("\n" + "=" * 50)
            print("所有测试完成！")
            return True
            
        finally:
            # 停止服务器
            self.stop_server()

if __name__ == "__main__":
    tester = KWDBTester()
    success = tester.run_all_tests()
    sys.exit(0 if success else 1)
